# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import os

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

from .types import State
from .nodes import (
    coordinator_node,
    planner_node,
    reporter_node_ogsm,
    reporter_node,
    research_team_node,
    researcher_node,
    coder_node,
    human_feedback_node,
    background_investigation_node, jd_generator_node, resume_matcher_node,
    question_generator_node, interview_evaluator_node, human_feedback2_node, reporter_checker_node,
    resume_matcher_loop_node, question_generator_loop_node, interview_evaluator_loop_node,
    test_start_node, test_start_node_ogsm,
    resume_matcher_summarize_node, interview_evaluator_summarize_node,
)
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.store.postgres import PostgresStore

POSTGRES_URL = os.getenv("POSTGRES_URL", "postgresql://postgres:postgres@47.108.188.77:54322/postgres")
DB_URI = POSTGRES_URL+"?sslmode=disable"


with (
    PostgresStore.from_conn_string(DB_URI) as store,
    PostgresSaver.from_conn_string(DB_URI) as checkpointer,
):

    def _build_base_graph():
        """Build and return the base state graph with all nodes and edges."""
        builder = StateGraph(State)
        builder.add_edge(START, "coordinator")
        builder.add_node("coordinator", coordinator_node)
        builder.add_node("background_investigator", background_investigation_node)
        builder.add_node("planner", planner_node)
        builder.add_node("reporter", reporter_node)
        builder.add_node("research_team", research_team_node)
        builder.add_node("researcher", researcher_node)
        builder.add_node("coder", coder_node)
        builder.add_node("human_feedback", human_feedback_node)
        builder.add_edge("reporter", END)
        return builder


    def _build_interview_graph():
        """Build and return the base state graph with all nodes and edges."""
        builder = StateGraph(State)
        builder.add_edge(START, "coordinator")
        builder.add_node("coordinator", coordinator_node)
        builder.add_node("background_investigator", background_investigation_node)
        builder.add_node("planner", planner_node)
        builder.add_node("research_team", research_team_node)
        builder.add_node("researcher", researcher_node)
        builder.add_node("coder", coder_node)
        builder.add_node("human_feedback", human_feedback_node)
        #
        builder.add_node("reporter", reporter_node)
        builder.add_node("reporter_checker", reporter_checker_node)
        builder.add_edge("reporter", "reporter_checker")
        # builder.add_edge(START, "test_start")
        builder.add_node("human_feedback2", human_feedback2_node)
        builder.add_node("test_start", test_start_node)
        builder.add_node("jd_generator", jd_generator_node)
        builder.add_node("resume_matcher_loop", resume_matcher_loop_node)
        builder.add_node("resume_matcher", resume_matcher_node)
        builder.add_node("resume_matcher_summarize", resume_matcher_summarize_node)
        builder.add_node("question_generator_loop", question_generator_loop_node)
        builder.add_node("question_generator", question_generator_node)
        builder.add_node("interview_evaluator_loop", interview_evaluator_loop_node)
        builder.add_node("interview_evaluator", interview_evaluator_node)
        builder.add_node("interview_evaluator_summarize", interview_evaluator_summarize_node)
        # builder.add_edge("interview_evaluator", END)

        return builder


    def _build_ogsm_graph():
        """Build and return the base state graph with all nodes and edges."""
        builder = StateGraph(State)
        # builder.add_edge(START, "test_start_node")
        builder.add_edge(START, "coordinator")

        # builder.add_node("test_start_node", test_start_node_ogsm)
        builder.add_node("coordinator", coordinator_node)
        builder.add_node("background_investigator", background_investigation_node)
        builder.add_node("planner", planner_node)
        builder.add_node("reporter", reporter_node_ogsm)
        builder.add_node("research_team", research_team_node)
        builder.add_node("researcher", researcher_node)
        builder.add_node("coder", coder_node)
        builder.add_node("human_feedback", human_feedback_node)
        builder.add_edge("reporter", END)
        return builder


    def build_graph_with_memory():
        """Build and return the agent workflow graph with memory."""
        # use persistent memory to save conversation history
        # 第一次执行要setup
        store.setup()
        checkpointer.setup()
        # memory = MemorySaver()

        # build state graph
        builder = _build_base_graph()
        # return builder.compile(checkpointer=memory)
        return builder.compile(
            checkpointer=checkpointer,
            store=store,
        )


    def build_graph_interview():
        """Build and return the agent workflow graph with memory."""
        # use persistent memory to save conversation history
        # 第一次执行要setup
        store.setup()
        checkpointer.setup()

        # memory = MemorySaver()
        # build state graph
        builder = _build_interview_graph()
        # return builder.compile(checkpointer=memory)
        return builder.compile(
            checkpointer=checkpointer,
            store=store,
        )


    def build_graph_ogsm():
        """Build and return the agent workflow graph with memory."""
        # use persistent memory to save conversation history
        #第一次执行要setup
        store.setup()
        checkpointer.setup()

        # memory = MemorySaver()

        # build state graph
        builder = _build_ogsm_graph()
        # return builder.compile(checkpointer=memory)
        return builder.compile(
            checkpointer=checkpointer,
            store=store,
        )


    def build_graph():
        """Build and return the agent workflow graph without memory."""
        builder = _build_base_graph()
        return builder.compile()


    graph = build_graph()
